import torch
import torch.nn as nn
import torch.nn.functional as F
from itertools import product

def fd_pad_diri_bc(x, pad=(1, 1, 1, 1), g = 0):
    x = F.pad(x, pad=pad, mode='constant', value=g)
    return x

def internal_conv_implicit(x):
    kernel = torch.tensor([[0, 0.25, 0], [0.25, 0, 0.25], [0, 0.25, 0]])
    kernel = kernel.type_as(x).view(1, 1, 3, 3).repeat(1, 1, 1, 1)
    with torch.no_grad():
        rhs = F.conv2d(x, kernel)
    return rhs

def internal_conv_explicit(x):
    kernel = torch.tensor([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
    kernel = kernel.type_as(x).view(1, 1, 3, 3).repeat(1, 1, 1, 1)
    rhs = F.conv2d(x, kernel)
    return rhs

def implicit_convRhs(h, k=1):
    h2 = h**2    
    force = lambda f: h2 * f/ (4 * k)

    conver = lambda x, f: internal_conv_implicit(fd_pad_diri_bc(x)) + force(f)[..., 1:-1, 1:-1]
    return conver

def explicit_convRhs(h, k=1):
    h2 = h**2    
    force = lambda f: f

    conver = lambda x, f: k * internal_conv_explicit(fd_pad_diri_bc(x)) / h2 \
                - force(f)[..., 1:-1, 1:-1]
    return  conver

